import autode.exceptions as ex

from autode.config import Config
from autode.log import logger
from autode.neb.ci import CINEB
from autode.transition_states.ts_guess import TSguess
from autode.utils import work_in

from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from autode.species.species import Species
    from autode.wrappers.methods import Method


def get_ts_guess_neb(
    reactant: "Species",
    product: "Species",
    method: "Method",
    name: str = "neb",
    n: int = 10,
) -> Optional[TSguess]:
    """
    Get a transition state guess using a nudged elastic band calculation. The
    geometry of the reactant is used as the fixed initial point and the final
    product geometry generated by driving a linear path to products, which is
    used as the initial guess for the NEB images

    ---------------------------------------------------------------------------
    Arguments:
        reactant (autode.species.Species):

        product (autode.species.Species):

        method (autode.wrappers.methods.Method):

        name (str):

        n (int): Number of images to use in the NEB

    Returns:
        (autode.transition_states.ts_guess.TSguess | None):
    """
    assert n is not None
    logger.info("Generating a TS guess using a nudged elastic band")

    neb = CINEB.from_end_points(reactant, product, num=n)

    @work_in(name)
    def calculate():
        return neb.calculate(method=method, n_cores=Config.n_cores)

    try:
        calculate()
    except ex.CalculationException:
        logger.error("NEB calculation failed")
        return None

    if neb.peak_species is None:
        logger.error(
            "Failed to find a peak in the NEB. Cannot create a TS guess"
        )
        return None

    return TSguess(
        atoms=neb.peak_species.atoms,
        reactant=reactant,
        product=product,
        name=name,
    )
